#include "gaussian.h"
#include "CholeskyDecomposition.h"
#include "LDL.h"
#include "DiagonalMatrix.h"
#include "backsubstitute.h"
#include "random.h"
#include <errors.h>
using namespace std;

namespace Math {

template <class T>
Gaussian<T>::Gaussian()
{
}

template <class T>
Gaussian<T>::Gaussian(int d)
{
  resize(d);
}

template <class T>
Gaussian<T>::Gaussian(const MatrixT& sigma, const VectorT& _mu)
{
  mu = _mu;
  setCovariance(sigma);
}

template <class T>
Gaussian<T>::Gaussian(const Gaussian& g)
  :L(g.L),mu(g.mu)
{
}

template <class T>
void Gaussian<T>::resize(int d)
{
  L.resize(d,d);
  mu.resize(d,Zero);
}

template <class T>
int Gaussian<T>::numDims() const
{
  return mu.n;
}

template <class T>
bool Gaussian<T>::setCovariance(const MatrixT& sigma)
{
  if(mu.n != 0 && sigma.m != mu.n) {
    FatalError("Invalid dimensions on covariance matrix");
    return false;
  }
  LDLDecomposition<T> chol;
  chol.set(sigma);
  DiagonalMatrixTemplate<T> D;
  chol.getD(D);
  chol.getL(L);
  for(int i=0;i<D.n;i++)
    D(i) = Sqrt(D(i));
  D.preMultiply(L,L);
  return true;
}

template <class T>
void Gaussian<T>::getCovariance(MatrixT& sigma) const
{
  sigma.mulTransposeB(L,L);
}

template <class T>
void Gaussian<T>::setMean(const VectorT& _mu)
{
  if(L.m != 0 && _mu.n != L.m) {
    FatalError("Invalid dimensions on mean vector");
  }
  mu = _mu;
}

template <class T>
T Gaussian<T>::probability(const VectorT& x) const
{
  //TODO: what if the determinant is zero?
  int d = numDims();
  T det=One;
  for(int i=0;i<d;i++) det*=L(i,i);
  T invc = Pow(2.0*Pi,0.5*T(d))*det;
  VectorT x_mu,y;
  x_mu.sub(x,mu);
  y.resize(L.n);
  LBackSubstitute(L,x_mu,y);
  return Exp(-Half*y.normSquared())/invc;
}

template <class T>
void Gaussian<T>::generate(VectorT& x) const
{
  //generate a normalized gaussian distributed variable y
  VectorT y(numDims());
  for(int i=0;i<y.n;i++) y(i) = RandGaussian();
  L.mul(y,x);
  x += mu;
}

template <class T>
void Gaussian<T>::setMarginalized(const Gaussian& g,int i)
{
  MatrixT cov;
  cov.mulTransposeB(g.L,g.L);
  VectorT mean1(g.mu.n-1);
  MatrixT cov11(cov.m-1,cov.n-1);
  for(int p=0;p<cov.m-1;p++) {
    int r=(p>=i?p+1:p);
    for(int q=0;q<cov.n-1;q++) {
      int c=(q>=i?q+1:q);
      cov11(p,q)=cov(r,c);
    }
    mean1(p)=g.mu(r);
  }
  mu = mean1;
  setCovariance(cov11);
}


template <class T>
void Gaussian<T>::setConditional(const Gaussian& g,Real xi,int i)
{
  MatrixT cov;
  cov.mulTransposeB(g.L,g.L);
  VectorT mean1(g.mu.n-1);
  MatrixT cov11(cov.m-1,cov.n-1);
  VectorT cov12(cov.m-1);
  for(int p=0;p<cov.m-1;p++) {
    int r=(p>=i?p+1:p);
    for(int q=0;q<cov.n-1;q++) {
      int c=(q>=i?q+1:q);
      cov11(p,q)=cov(r,c);
    }
    cov12(p)=cov(r,i);
    mean1(p)=g.mu(r);
  }

  Real invC22=1.0/cov(i,i);
  mu = mean1;
  mu.madd(cov12,(xi-g.mu(i))*invC22);
  for(int p=0;p<cov11.m;p++)
    for(int q=0;q<cov11.n;q++)
      cov11(p,q) -= cov12(p)*cov12(q)*invC22;
  setCovariance(cov11);
}

template class Gaussian<float>;
template class Gaussian<double>;

} //namespace Math

